import streamlit as st
import os
from dotenv import load_dotenv
from llm_client import LLMClient
from emotion_engine import EmotionEngine
from prompts import SYSTEM_PROMPT

# 加载环境变量
load_dotenv()

# 页面配置
st.set_page_config(
    page_title="有感情的AI助手",
    page_icon="💭",
    layout="wide",
    initial_sidebar_state="expanded"
)

# 自定义CSS
st.markdown("""
    <style>
    .chat-message {
        padding: 1.5rem;
        border-radius: 0.5rem;
        margin-bottom: 1rem;
        display: flex;
        gap: 1rem;
    }
    .user-message {
        background-color: #e3f2fd;
        border-left: 4px solid #2196f3;
    }
    .assistant-message {
        background-color: #f3e5f5;
        border-left: 4px solid #9c27b0;
    }
    .emotion-badge {
        display: inline-block;
        padding: 0.25rem 0.75rem;
        border-radius: 1rem;
        font-size: 0.85rem;
        font-weight: bold;
        margin-top: 0.5rem;
    }
    .emotion-happy { background-color: #fff9c4; color: #f57f17; }
    .emotion-sad { background-color: #b3e5fc; color: #01579b; }
    .emotion-angry { background-color: #ffccbc; color: #d84315; }
    .emotion-calm { background-color: #c8e6c9; color: #1b5e20; }
    .emotion-confused { background-color: #ffe0b2; color: #e65100; }
    </style>
""", unsafe_allow_html=True)

# 初始化会话状态
if "messages" not in st.session_state:
    st.session_state.messages = []
if "llm_client" not in st.session_state:
    st.session_state.llm_client = LLMClient()
if "emotion_engine" not in st.session_state:
    st.session_state.emotion_engine = EmotionEngine()

# 侧边栏配置
with st.sidebar:
    st.title("⚙️ 配置")
    
    # API配置
    api_provider = st.selectbox(
        "选择LLM服务商",
        ["OpenAI", "Claude", "Ollama", "本地模型"]
    )
    
    if api_provider == "OpenAI":
        api_key = st.text_input("OpenAI API Key", type="password", key="openai_key")
        model = st.selectbox("模型选择", ["gpt-4", "gpt-3.5-turbo"])
    elif api_provider == "Claude":
        api_key = st.text_input("Claude API Key", type="password", key="claude_key")
        model = st.selectbox("模型选择", ["claude-3-opus", "claude-3-sonnet"])
    else:
        api_key = "local"
        model = st.text_input("本地模型名称", "llama2")
    
    # 情感配置
    st.divider()
    st.subheader("😊 情感设置")
    
    personality = st.selectbox(
        "选择AI的性格",
        ["温暖友善", "专业理性", "幽默风趣", "同情共鸣", "充满活力"]
    )
    
    emotion_intensity = st.slider(
        "情感强度",
        min_value=0.0,
        max_value=1.0,
        value=0.7,
        step=0.1
    )
    
    temperature = st.slider(
        "创意度（Temperature）",
        min_value=0.0,
        max_value=1.0,
        value=0.8,
        step=0.1
    )
    
    # 清空对话历史
    if st.button("🗑️ 清空对话", use_container_width=True):
        st.session_state.messages = []
        st.rerun()
    
    # 显示模型信息
    st.divider()
    st.info(f"当前模型: {model}\n情感强度: {emotion_intensity}\n性格: {personality}")

# 主页面标题
col1, col2 = st.columns([3, 1])
with col1:
    st.title("💭 有感情的AI助手")
    st.markdown("与一个具有真实情感的AI进行对话")

# 显示聊天历史
st.subheader("对话历史")
for message in st.session_state.messages:
    if message["role"] == "user":
        st.markdown(f"""
        <div class="chat-message user-message">
            <div style="flex: 1;">
                <strong>👤 你:</strong>
                <div style="margin-top: 0.5rem;">{message["content"]}</div>
            </div>
        </div>
        """, unsafe_allow_html=True)
    else:
        emotion = message.get("emotion", "calm")
        emotion_text = message.get("emotion_text", "平静")
        emotion_class = f"emotion-{emotion}"
        
        st.markdown(f"""
        <div class="chat-message assistant-message">
            <div style="flex: 1;">
                <strong>🤖 AI助手:</strong>
                <div style="margin-top: 0.5rem;">{message["content"]}</div>
                <span class="emotion-badge {emotion_class}">💗 {emotion_text}</span>
            </div>
        </div>
        """, unsafe_allow_html=True)

# 输入区域
st.divider()
st.subheader("💬 输入你的消息")

# 使用输入框
user_input = st.text_area(
    "请输入你的消息:",
    placeholder="与AI进行对话...",
    height=100,
    key="user_input"
)

col1, col2, col3 = st.columns([2, 1, 1])

with col1:
    send_button = st.button("📤 发送", use_container_width=True, type="primary")

with col2:
    example_button = st.button("💡 示例", use_container_width=True)

with col3:
    import_button = st.button("📥 导入对话", use_container_width=True)

# 处理示例按钮
if example_button:
    examples = [
        "我今天感到很难过，可以陪我聊天吗？",
        "你觉得人生的意义是什么？",
        "给我讲一个励志的故事"
    ]
    selected_example = st.selectbox("选择一个示例:", examples)
    user_input = selected_example

# 处理发送消息
if send_button and user_input:
    # 添加用户消息到历史
    st.session_state.messages.append({
        "role": "user",
        "content": user_input
    })
    
    # 显示加载状态
    with st.spinner("AI正在思考..."):
        try:
            # 构建系统提示
            system_message = SYSTEM_PROMPT.format(
                personality=personality,
                emotion_intensity=emotion_intensity
            )
            
            # 调用LLM
            llm_response = st.session_state.llm_client.chat(
                messages=st.session_state.messages,
                system_prompt=system_message,
                model=model,
                temperature=temperature,
                api_key=api_key,
                provider=api_provider
            )
            
            # 分析情感
            emotion, emotion_text = st.session_state.emotion_engine.analyze(
                user_input,
                llm_response
            )
            
            # 添加AI响应到历史
            st.session_state.messages.append({
                "role": "assistant",
                "content": llm_response,
                "emotion": emotion,
                "emotion_text": emotion_text
            })
            
            st.rerun()
            
        except Exception as e:
            st.error(f"发生错误: {str(e)}")
            st.info("请检查API配置和网络连接")

# 页脚
st.divider()
st.caption("🌟 这是一个有感情的AI助手，旨在提供温暖、理解和支持。")
